import chromadb
from langchain_chroma import Chroma
from langchain_huggingface import HuggingFaceEmbeddings


class LocalEmbedding:
    def __init__(self, model_name="all-MiniLM-L6-v2"):
        self.embedding_model = HuggingFaceEmbeddings(
            model_name=model_name,  # 模型名称
            model_kwargs={"device": "cpu"},  # 设备
            encode_kwargs={"normalize_embedding": True},  # 归一化
        )

    def get(self):
        return self.embedding_model


# 插入测试数据到Chroma数据库的函数
def insert_test_data():
    # 打印初始化Chroma数据库日志
    print("[DB] 初始化Chroma数据库 product_db")
    # 创建Chroma数据库对象
    chromaClient = Chroma(
        "product_db",  # 数据库名称
        LocalEmbedding().get(),  # 嵌入模型
        client=chromadb.PersistentClient(path="./chroma_db"),  # 数据库客户端
    )
    # 定义测试数据列表，每条包含内容和元数据
    test_docs = [
        {
            "content": "如需申请手机保修服务，请携带购机发票和保修卡前往品牌授权售后服务中心，工程师检测后符合保修政策即可免费维修。",
            "meta": {"category": "手机保修"},
        },
        {
            "content": "笔记本电脑电池出现鼓包属于安全隐患，请立即停止使用并联系品牌售后服务中心进行更换，切勿自行拆卸或继续充电。",
            "meta": {"category": "电脑电池"},
        },
        {
            "content": "电子产品自购买之日起享受7天无理由退货，15天内可换货，一年内享受免费保修服务，具体以品牌政策为准。",
            "meta": {"category": "售后政策"},
        },
        {
            "content": "如遇产品无法开机、屏幕碎裂等问题，请及时联系官方售后或客服，部分问题可能不在保修范围内。",
            "meta": {"category": "常见问题"},
        },
    ]

    # 遍历每条测试数据，插入到Chroma数据库
    for i, doc in enumerate(test_docs):
        # 打印插入测试数据日志
        print(f"[DB] 插入测试数据: {doc['meta']['category']}")
        # 调用add_texts方法插入文本、元数据和ID
        chromaClient.add_texts(
            texts=[doc["content"]], metadatas=[doc["meta"]], ids=[f"test_doc_{i}"]
        )
    # 打印测试数据插入完成日志
    print("[DB] 测试数据已插入 product_db。")


if __name__ == "__main__":
    insert_test_data()
